package cn.iocoder.yudao.framework.datasource.config;

import cn.iocoder.yudao.framework.datasource.core.DynamicRoutingDataSource;
import com.alibaba.druid.pool.DruidDataSource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.boot.autoconfigure.AutoConfigureAfter;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.boot.autoconfigure.jdbc.DataSourceAutoConfiguration;
import org.springframework.boot.context.properties.EnableConfigurationProperties;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.Primary;
import org.springframework.core.env.Environment;

import javax.sql.DataSource;
import java.util.HashMap;
import java.util.Map;

/**
 * 芋道动态数据源自动配置类
 * 替换原有的dynamic-datasource实现，使用自定义的DynamicRoutingDataSource
 *
 * @author 芋道源码
 */
@Slf4j
@Configuration
@EnableConfigurationProperties(DynamicDataSourceProperties.class)
@AutoConfigureAfter(DataSourceAutoConfiguration.class)
@ConditionalOnProperty(prefix = "spring.datasource.dynamic", name = "enabled", havingValue = "true", matchIfMissing = true)
public class YudaoDynamicDataSourceAutoConfiguration {

    /**
     * 创建动态路由数据源
     */
    @Bean
    @Primary
    public DynamicRoutingDataSource dynamicDataSource(Environment environment, 
                                                     DynamicDataSourceProperties properties) {
        DynamicRoutingDataSource dynamicRoutingDataSource = new DynamicRoutingDataSource();
        
        // 构建数据源映射
        Map<Object, Object> targetDataSources = new HashMap<>();
        
        // 添加主数据源
        DataSource masterDataSource = createDataSource(environment, "master", properties);
        targetDataSources.put("master", masterDataSource);
        
        // 添加从数据源（如果配置了的话）
        if (hasSlaveConfig(environment)) {
            DataSource slaveDataSource = createDataSource(environment, "slave", properties);
            targetDataSources.put("slave", slaveDataSource);
        }
        
        // 设置目标数据源
        dynamicRoutingDataSource.setTargetDataSources(targetDataSources);
        dynamicRoutingDataSource.setDefaultTargetDataSource(masterDataSource);
        
        log.info("Dynamic DataSource 初始化完成，包含数据源: {}", targetDataSources.keySet());
        
        return dynamicRoutingDataSource;
    }

    /**
     * 创建Druid数据源
     */
    private DataSource createDataSource(Environment environment, String name, 
                                      DynamicDataSourceProperties properties) {
        DruidDataSource dataSource = new DruidDataSource();
        
        // 基本连接信息
        String prefix = "spring.datasource.dynamic.datasource." + name + ".";
        dataSource.setUrl(environment.getProperty(prefix + "url"));
        dataSource.setUsername(environment.getProperty(prefix + "username"));
        dataSource.setPassword(environment.getProperty(prefix + "password"));
        dataSource.setDriverClassName(environment.getProperty(prefix + "driver-class-name", 
                "com.mysql.cj.jdbc.Driver"));
        
        // Druid连接池配置
        DynamicDataSourceProperties.DruidConfig druidConfig = properties.getDruid();
        if (druidConfig != null) {
            dataSource.setInitialSize(druidConfig.getInitialSize());
            dataSource.setMinIdle(druidConfig.getMinIdle());
            dataSource.setMaxActive(druidConfig.getMaxActive());
            dataSource.setMaxWait(druidConfig.getMaxWait());
            dataSource.setTimeBetweenEvictionRunsMillis(druidConfig.getTimeBetweenEvictionRunsMillis());
            dataSource.setMinEvictableIdleTimeMillis(druidConfig.getMinEvictableIdleTimeMillis());
            dataSource.setMaxEvictableIdleTimeMillis(druidConfig.getMaxEvictableIdleTimeMillis());
            dataSource.setValidationQuery(druidConfig.getValidationQuery());
            dataSource.setTestWhileIdle(druidConfig.getTestWhileIdle());
            dataSource.setTestOnBorrow(druidConfig.getTestOnBorrow());
            dataSource.setTestOnReturn(druidConfig.getTestOnReturn());
            dataSource.setPoolPreparedStatements(druidConfig.getPoolPreparedStatements());
            dataSource.setMaxPoolPreparedStatementPerConnectionSize(
                    druidConfig.getMaxPoolPreparedStatementPerConnectionSize());
        }
        
        try {
            dataSource.init();
            log.info("初始化数据源 {} 成功", name);
        } catch (Exception e) {
            log.error("初始化数据源 {} 失败", name, e);
            throw new RuntimeException("初始化数据源失败: " + name, e);
        }
        
        return dataSource;
    }

    /**
     * 检查是否配置了从数据源
     */
    private boolean hasSlaveConfig(Environment environment) {
        String slaveUrl = environment.getProperty("spring.datasource.dynamic.datasource.slave.url");
        return slaveUrl != null && !slaveUrl.trim().isEmpty();
    }
}